import os
import shutil
import tempfile
import zipfile
import logging
from typing import List, Optional
from fastapi import FastAPI, UploadFile, HTTPException, Body
from fastapi.responses import JSONResponse, FileResponse
from fastapi.middleware.cors import CORSMiddleware
from pdf2image import convert_from_path
import pytesseract
from PIL import Image
from deepseek import DeepSeekAPI
import markdown
from fastapi import APIRouter
from pydantic import BaseModel
from prompt_manager import PromptManager

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('app.log')
    ]
)
logger = logging.getLogger(__name__)

app = FastAPI()
router = APIRouter(prefix="/api")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3003"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.include_router(router)

# 创建 PromptManager 实例
prompt_manager = PromptManager()

# 创建 DeepSeekAPI 实例
deepseek_api = DeepSeekAPI()

# 文件处理状态字典
file_processing_status = {}

# 存储OCR结果的字典
ocr_results = {}

# 定义请求模型
class ProcessRequestModel(BaseModel):
    temp_dir: str
    prompt_type: Optional[str] = "translate"  # 添加可选的 prompt 类型字段
    target_language: Optional[str] = "en"  # 添加目标语言字段
    custom_prompt: Optional[str] = None  # 添加自定义prompt字段
    
    class Config:
        # Pydantic V1 兼容配置
        schema_extra = {
            "example": {
                "temp_dir": "/path/to/temp/dir",
                "prompt_type": "translate",
                "target_language": "en",
                "custom_prompt": "请将文本翻译成英文，并保持专业术语的准确性"
            }
        }

@app.get("/")
async def root():
    return {"message": "Welcome to PDF OCR & Translation Tool"}

@router.get("/status")
async def api_status():
    return {"status": "API is running"}

# 第1步：上传PDF压缩包
@router.post("/upload")
async def upload_file(file: UploadFile):
    logger.info(f"Received upload request for file: {file.filename}")
    try:
        # 验证文件名和类型
        if not file.filename or not file.filename.endswith('.zip'):
            logger.error(f"Invalid file type or name: {file.filename}")
            raise HTTPException(status_code=400, detail="Only ZIP files are allowed")

        # 验证文件大小（限制为100MB）
        logger.info("Starting file size validation")
        file_size = 0
        chunk_size = 1024 * 1024  # 1MB
        content = await file.read(chunk_size)
        while content:
            file_size += len(content)
            if file_size > 100 * 1024 * 1024:  # 100MB
                logger.error(f"File too large: {file.filename}, size: {file_size} bytes")
                raise HTTPException(status_code=400, detail="File size exceeds 100MB limit")
            content = await file.read(chunk_size)
        await file.seek(0)
        logger.info(f"File size validation passed: {file_size} bytes")

        # 创建并清理目录
        logger.info("Setting up directories")
        upload_dir = os.path.join(os.path.dirname(__file__), 'upload')
        extract_dir = os.path.join(upload_dir, 'extracted')
        output_dir = os.path.join(upload_dir, 'output')
        ocr_dir = os.path.join(upload_dir, 'ocr')
        ocr_img_dir = os.path.join(ocr_dir, 'img')
        ocr_text_dir = os.path.join(ocr_dir, 'text')

        # 确保目录存在
        try:
            os.makedirs(upload_dir, exist_ok=True)
            os.makedirs(extract_dir, exist_ok=True)
            os.makedirs(output_dir, exist_ok=True)
            os.makedirs(ocr_dir, exist_ok=True)
            os.makedirs(ocr_img_dir, exist_ok=True)
            os.makedirs(ocr_text_dir, exist_ok=True)
            logger.info("Directories created successfully")
        except Exception as e:
            logger.error(f"Failed to create directories: {str(e)}")
            raise HTTPException(status_code=500, detail="Failed to create required directories")

        # 清理之前的文件
        for dir_path in [extract_dir, output_dir, ocr_img_dir, ocr_text_dir]:
            for item in os.listdir(dir_path):
                item_path = os.path.join(dir_path, item)
                try:
                    if os.path.isfile(item_path):
                        os.unlink(item_path)
                    elif os.path.isdir(item_path):
                        shutil.rmtree(item_path)
                except Exception as e:
                    logger.warning(f"Error cleaning up {item_path}: {str(e)}")

        # 保存上传的文件
        zip_path = os.path.join(upload_dir, file.filename)
        logger.info(f"Attempting to save uploaded file to: {zip_path}")
        try:
            with open(zip_path, 'wb') as f:
                shutil.copyfileobj(file.file, f)
            if not os.path.exists(zip_path):
                logger.error("File was not saved successfully")
                raise HTTPException(status_code=500, detail="File save verification failed")
            logger.info(f"File saved successfully, size: {os.path.getsize(zip_path)} bytes")
        except Exception as e:
            logger.error(f"Failed to save file: {str(e)}, type: {type(e).__name__}")
            raise HTTPException(status_code=500, detail=f"Failed to save uploaded file: {str(e)}")

        # 第2步：解压文件，使用zipfile处理中文编码
        logger.info("Starting ZIP file extraction")
        try:
            with zipfile.ZipFile(zip_path, 'r') as zip_ref:
                logger.info(f"ZIP file opened successfully, contains {len(zip_ref.filelist)} files")
                for file_info in zip_ref.filelist:
                    try:
                        original_filename = file_info.filename
                        logger.info(f"Processing ZIP entry: {original_filename}")
                        # 使用utf-8编码处理文件名
                        try:
                            # 先将cp437编码的文件名转换为bytes，再用utf-8解码
                            file_info.filename = file_info.filename.encode('cp437').decode('utf-8')
                            logger.info(f"Successfully decoded filename: {file_info.filename}")
                        except UnicodeDecodeError:
                            # 如果utf-8解码失败，尝试使用原始文件名
                            logger.warning(f"Failed to decode filename, using original: {original_filename}")
                            file_info.filename = original_filename
                        
                        # 检查文件名是否合法
                        if not os.path.basename(file_info.filename):
                            continue
                            
                        zip_ref.extract(file_info, extract_dir)
                    except Exception as e:
                        logger.error(f"Error extracting {file_info.filename}: {str(e)}")
                        continue

            # 获取所有PDF文件
            pdf_files = [f for f in os.listdir(extract_dir) if f.lower().endswith('.pdf')]
            
            if not pdf_files:
                raise HTTPException(status_code=400, detail="No PDF files found in the ZIP archive")
            
            # 确保目录权限正确
            os.chmod(upload_dir, 0o755)
            os.chmod(extract_dir, 0o755)
            os.chmod(output_dir, 0o755)
            os.chmod(ocr_dir, 0o755)
            os.chmod(ocr_img_dir, 0o755)
            os.chmod(ocr_text_dir, 0o755)
            
            response_data = {
                "files": pdf_files,
                "temp_dir": upload_dir
            }
            logger.info(f"Upload successful. Response: {response_data}")
            return JSONResponse(response_data)
        
        except zipfile.BadZipFile:
            logger.error("Invalid ZIP file format")
            raise HTTPException(status_code=400, detail="Invalid ZIP file")
        except Exception as e:
            logger.error(f"Upload error: {str(e)}")
            raise HTTPException(status_code=500, detail="Internal server error during file processing")
            
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

# 第3步：OCR处理PDF获取图片和文本
@router.post("/ocr/{filename}")
async def ocr_file(filename: str, request: ProcessRequestModel):
    from urllib.parse import unquote
    decoded_filename = unquote(filename)
    temp_dir = request.temp_dir
    
    logger.info(f"Received OCR request - Filename: {decoded_filename}, Temp directory: {temp_dir}")
    try:
        # 初始化处理状态
        file_processing_status[decoded_filename] = {
            "status": "processing",
            "progress": 0,
            "message": "开始OCR处理"
        }
        logger.info(f"Initialized OCR processing status: {file_processing_status[decoded_filename]}")

        # 获取PDF文件路径
        pdf_path = os.path.join(temp_dir, 'extracted', decoded_filename)
        if not os.path.exists(pdf_path):
            error_msg = f"PDF file not found at path: {pdf_path}"
            logger.error(error_msg)
            file_processing_status[decoded_filename] = {"status": "error", "message": "文件未找到"}
            raise HTTPException(status_code=404, detail="File not found")
        
        # 创建OCR图片和文本目录
        ocr_img_dir = os.path.join(temp_dir, 'ocr', 'img')
        ocr_text_dir = os.path.join(temp_dir, 'ocr', 'text')
        os.makedirs(ocr_img_dir, exist_ok=True)
        os.makedirs(ocr_text_dir, exist_ok=True)
        logger.info(f"Created OCR directories: {ocr_img_dir}, {ocr_text_dir}")
        
        # 将PDF转换为图片
        file_processing_status[decoded_filename].update({"progress": 20, "message": "正在转换PDF为图片"})
        logger.info(f"Converting PDF to images: {file_processing_status[decoded_filename]}")
        images = convert_from_path(pdf_path)
        
        # 保存图片并OCR处理
        text_content = []
        image_paths = []
        total_images = len(images)
        
        base_filename = os.path.splitext(decoded_filename)[0]
        
        for idx, image in enumerate(images, 1):
            # 保存图片
            img_filename = f"{base_filename}_page_{idx}.png"
            img_path = os.path.join(ocr_img_dir, img_filename)
            image.save(img_path, 'PNG')
            image_paths.append(img_path)
            logger.info(f"Saved image to: {img_path}")
            
            # OCR处理
            status_update = {
                "progress": 20 + int(60 * (idx / total_images)),
                "message": f"正在进行OCR处理 ({idx}/{total_images})"
            }
            file_processing_status[decoded_filename].update(status_update)
            logger.info(f"OCR processing status: {status_update}")
            
            text = pytesseract.image_to_string(image, lang='chi_sim+eng', config='--oem 1 --psm 3')
            text_content.append(text)
            
            # 保存每页的OCR文本
            text_filename = f"{base_filename}_page_{idx}.txt"
            text_path = os.path.join(ocr_text_dir, text_filename)
            with open(text_path, 'w', encoding='utf-8') as f:
                f.write(text)
        
        # 合并文本并确保使用UTF-8编码
        full_text = '\n'.join(text_content).encode('utf-8').decode('utf-8')
        
        # 保存完整的OCR文本
        full_text_filename = f"{base_filename}_full.txt"
        full_text_path = os.path.join(ocr_text_dir, full_text_filename)
        with open(full_text_path, 'w', encoding='utf-8') as f:
            f.write(full_text)
        
        # 存储OCR结果
        ocr_results[decoded_filename] = {
            "text": full_text,
            "image_paths": image_paths,
            "text_path": full_text_path
        }
        
        # 更新最终状态
        final_status = {
            "status": "success",
            "progress": 100,
            "message": "OCR处理完成",
            "file_info": {
                "filename": decoded_filename,
                "pages": total_images,
                "text_path": full_text_path,
                "image_count": len(image_paths)
            }
        }
        file_processing_status[decoded_filename] = final_status
        logger.info(f"OCR processing completed successfully: {final_status}")
        
        return JSONResponse(final_status)
    
    except Exception as e:
        error_msg = f"Error in OCR processing for file {decoded_filename}: {str(e)}"
        logger.error(error_msg)
        error_status = {
            "status": "error",
            "progress": 0,
            "message": str(e)
        }
        file_processing_status[decoded_filename] = error_status
        logger.error(f"Error status: {error_status}")
        raise HTTPException(status_code=500, detail=str(e))

# 第4步：获取可用的prompt类型
@router.get("/prompts")
async def get_available_prompts():
    """获取所有可用的 prompt 模板"""
    try:
        available_prompts = prompt_manager.list_available_prompts()
        return {
            "status": "success",
            "data": {
                "prompts": available_prompts
            }
        }
    except Exception as e:
        logger.error(f"获取 prompt 模板列表时出错: {str(e)}")
        return JSONResponse(
            status_code=500,
            content={
                "status": "error",
                "message": f"获取 prompt 模板列表失败: {str(e)}",
                "error_code": "PROMPT_LIST_ERROR"
            }
        )

# 获取处理状态
@router.get("/status/{filename}")
async def get_processing_status(filename: str):
    logger.info(f"Received status check request for file: {filename}")
    if filename not in file_processing_status:
        response = {"status": "not_found"}
        logger.info(f"Status check response: {response}")
        return JSONResponse(response)
    response = file_processing_status[filename]
    logger.info(f"Status check response: {response}")
    return JSONResponse(response)

# 第5步：提交OCR文本给大模型处理
@router.post("/process/{filename}")
async def process_file(filename: str, request: ProcessRequestModel):
    from urllib.parse import unquote
    decoded_filename = unquote(filename)
    temp_dir = request.temp_dir
    prompt_type = request.prompt_type or "translate"  # 默认使用翻译 prompt
    target_language = request.target_language or "en"  # 默认目标语言为英语
    custom_prompt = request.custom_prompt  # 获取自定义prompt
    
    logger.info(f"Received process request - Filename: {decoded_filename}, Temp directory: {temp_dir}, Prompt type: {prompt_type}, Target language: {target_language}, Custom prompt: {custom_prompt}")
    try:
        # 检查文件是否已经过OCR处理
        if decoded_filename not in ocr_results:
            error_msg = f"File {decoded_filename} has not been OCR processed yet"
            logger.error(error_msg)
            raise HTTPException(status_code=400, detail="Please perform OCR first")
        
        # 初始化处理状态
        file_processing_status[decoded_filename] = {
            "status": "processing",
            "progress": 0,
            "message": "开始处理文本"
        }
        logger.info(f"Initialized text processing status: {file_processing_status[decoded_filename]}")
        
        # 获取OCR结果
        ocr_result = ocr_results[decoded_filename]
        full_text = ocr_result["text"]
        
        # 更新状态
        file_processing_status[decoded_filename].update({"progress": 30, "message": f"正在使用 {prompt_type} 模式处理"})
        logger.info(f"Processing with {prompt_type} mode: {file_processing_status[decoded_filename]}")
        
        # 使用DeepSeek API进行处理
        processed_text = deepseek_api.process(
            text=full_text,
            prompt_type=prompt_type,
            target_lang=target_language,
            custom_prompt=custom_prompt
        )
        
        # 转换为Markdown格式并保存
        file_processing_status[decoded_filename].update({"progress": 80, "message": "正在保存结果"})
        logger.info(f"Saving results status: {file_processing_status[decoded_filename]}")
        
        base_filename = os.path.splitext(decoded_filename)[0]
        md_content = f"# {base_filename}\n\n{processed_text}"
        
        # 保存Markdown文件
        output_path = os.path.join(temp_dir, 'output', f"{base_filename}_{prompt_type}.md")
        with open(output_path, 'w', encoding='utf-8') as f:
            f.write(md_content)
        
        # 更新最终状态
        final_status = {
            "status": "success",
            "progress": 100,
            "message": "处理完成",
            "output_file": output_path,
            "filename": f"{base_filename}_{prompt_type}.md"
        }
        file_processing_status[decoded_filename] = final_status
        logger.info(f"Processing completed successfully: {final_status}")
        
        return JSONResponse(final_status)
    
    except Exception as e:
        error_msg = f"Error processing file {decoded_filename}: {str(e)}"
        logger.error(error_msg)
        error_status = {
            "status": "error",
            "progress": 0,
            "message": str(e)
        }
        file_processing_status[decoded_filename] = error_status
        logger.error(f"Error status: {error_status}")
        raise HTTPException(status_code=500, detail=str(e))

# 第6步：下载处理结果的markdown文档
@router.get("/download/{filename}")
async def download_file(filename: str, temp_dir: str = None):
    logger.info(f"Received download request for file: {filename}")
    try:
        # 获取文件路径
        if temp_dir:
            # 如果提供了temp_dir，直接使用完整路径
            file_path = os.path.join(temp_dir, 'ocr', 'text', filename)
        else:
            # 否则使用默认的output目录
            file_path = os.path.join(os.path.dirname(__file__), 'upload', 'output', filename)
        
        if not os.path.exists(file_path):
            error_msg = f"File not found at path: {file_path}"
            logger.error(error_msg)
            raise HTTPException(status_code=404, detail="File not found")
        
        logger.info(f"Serving file for download: {file_path}")
        # 根据文件扩展名设置正确的media_type
        media_type = "text/plain" if filename.endswith('.txt') else "text/markdown"
        return FileResponse(
            path=file_path, 
            filename=filename,
            media_type=media_type
        )
    
    except Exception as e:
        error_msg = f"Error downloading file {filename}: {str(e)}"
        logger.error(error_msg)
        raise HTTPException(status_code=500, detail=str(e))

app.include_router(router)

if __name__ == "__main__":
    logger.info("Starting server on 0.0.0.0:5063")
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=5063)